from NN import ann
import torch
import time

from verify_mip_relu import milp_verify

# import gurobipy as gp
# from gurobipy import GRB


def verify_relu(model, piece_len):
    cex_init, res_init = milp_verify.MILP_opt_initCon(model)
    cex_unsafe, res_unsafe = milp_verify.MILP_opt_unsafeCon(model)
    cex_domain, res_domain = milp_verify.MILP_opt_thirdCon(model, piece_len)
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = ann.gen_nn()
    model.load_state_dict(torch.load('../ann-model/ann-trained_eg2_relu.pt'), strict=True)

    time_start_verify = time.time()

    verify_relu(model)

    time_end_verify = time.time()
    print("verifing totally costs:", time_end_verify - time_start_verify)

